{
    "cells": [
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "52295407",
            "metadata": {},
            "outputs": [],
            "source": [
                "import logging\n",
                "import sys\n",
                "\n",
                "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
                "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "c5d167a5-81f8-4d2c-b42f-0a190577132f",
            "metadata": {},
            "outputs": [],
            "source": [
                "# My OpenAI Key\n",
                "import os\n",
                "\n",
                "os.environ[\"OPENAI_API_KEY\"] = \"INSERT OPENAI KEY\""
            ]
        },
        {
            "cell_type": "markdown",
            "id": "575750cc-479f-4b1f-b93f-4b00ed756d52",
            "metadata": {},
            "source": [
                "## Wikipedia Reader + Keyword Table"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 36,
            "id": "5f60348e-731d-4a95-bae2-426e184a914e",
            "metadata": {},
            "outputs": [],
            "source": [
                "from llama_index import KeywordTableIndex, WikipediaReader"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 37,
            "id": "952c4659-7fbb-447e-8caf-06916412cc37",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "page: Covid-19\n"
                    ]
                }
            ],
            "source": [
                "wiki_docs = WikipediaReader().load_data(pages=[\"Covid-19\"])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "3be202db-a4c7-41d2-ba7d-446d1f934830",
            "metadata": {},
            "outputs": [],
            "source": [
                "index = KeywordTableIndex.from_documents(wiki_docs)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 46,
            "id": "28d7163e-f26f-4ad8-89d5-9cb7662c4d9c",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "> Starting query: Which country included tocilizumab in treatment for covid-19?\n",
                        "Extracted keywords: ['tocilizumab', 'treatment', 'covid-19', 'covid', '19']\n",
                        "> Querying with idx: 1105763466456338724: of age or older weighing at least 40 kilograms ...\n",
                        "> Querying with idx: 2820318727532393752: Coronavirus disease 2019 (COVID-19) is a contag...\n",
                        "> Querying with idx: 897499143815831368: if the mask includes an exhalation valve, a wea...\n",
                        "> Querying with idx: 8628144746434065339: pulmonary fibrosis, cystic fibrosis. Evidence s...\n"
                    ]
                },
                {
                    "data": {
                        "text/plain": [
                            "'\\n\\nChina'"
                        ]
                    },
                    "execution_count": 46,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "# GPT doesn't find the corresponding evidence in the leaf node, but still gives the correct answer\n",
                "# set Logging to DEBUG for more detailed outputs\n",
                "query_engine = index.as_query_engine()\n",
                "query_engine.query(\"Which country included tocilizumab in treatment for covid-19?\")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "addb0c4d-f1ae-40c1-8b69-5a989609672f",
            "metadata": {},
            "source": [
                "## Wikipedia Reader + List"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "a0fc24e1-eca5-4267-a962-f7fe0fc5c7df",
            "metadata": {},
            "outputs": [],
            "source": [
                "from llama_index import SummaryIndex, WikipediaReader"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "id": "872a651a-ca4a-43e2-8b29-e4f667f9d3c5",
            "metadata": {},
            "outputs": [],
            "source": [
                "wiki_docs = WikipediaReader().load_data(pages=[\"Covid-19\"])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "id": "37e85af0-b1c3-4c18-b239-6e32a7acf8d6",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "> Adding chunk: Coronavirus disease 2019 (COVID-19) is a contag...\n",
                        "> Adding chunk: people with COVID‑19 and acute respiratory dist...\n",
                        "> Adding chunk: encourage or mandate the use of face masks or c...\n",
                        "> Adding chunk: have elevated liver enzymes, reflecting liver i...\n",
                        "> Adding chunk: insofar as their drug use may have caused lung ...\n",
                        "> Adding chunk: treatment of mild-to-moderate COVID‑19 in adult...\n"
                    ]
                }
            ],
            "source": [
                "index = SummaryIndex.from_documents(wiki_docs)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 13,
            "id": "ec0119ef-786e-40ea-89af-f1ca0ad26de6",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "> Starting query: Which country included tocilizumab in treatment for covid-19?\n"
                    ]
                }
            ],
            "source": [
                "# set Logging to DEBUG for more detailed outputs\n",
                "# with keyword lookup\n",
                "from llama_index.indices.postprocessor import KeywordNodePostprocessor\n",
                "\n",
                "\n",
                "query_engine = index.as_query_engine(\n",
                "    node_postprocessors=[KeywordNodePostprocessor(required_keywords=[\"tocilizumab\"])]\n",
                ")\n",
                "response = query_engine.query(\n",
                "    \"Which country included tocilizumab in treatment for covid-19?\",\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 18,
            "id": "b4087a84-0939-444f-93f2-a1a7aa32db3f",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "'China'"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                }
            ],
            "source": [
                "display(response.strip())"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 19,
            "id": "fb155bc7-cb50-47b6-b92b-895852c2d8f4",
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "> Starting query: Which country included tocilizumab in treatment for covid-19?\n"
                    ]
                }
            ],
            "source": [
                "# set Logging to DEBUG for more detailed outputs\n",
                "# without keyword lookup\n",
                "query_engine = index.as_query_engine()\n",
                "response = query_engine.query(\n",
                "    \"Which country included tocilizumab in treatment for covid-19?\"\n",
                ")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 20,
            "id": "5b45c07a-4e76-4a45-86b6-6b2df1ef4f7b",
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "'There is no definite answer to this question as different countries have different treatment methods for covid-19. However, according to the context information, it is known that the virus SARS-CoV-2 can cause severe damage to various organs in the human body by inducing systemic inflammation. Therefore, it is possible that tocilizumab, which is a drug that inhibits the virus, may be included in treatment for covid-19 in some countries in order to prevent or reduce the severity of a cytokine storm. Additionally, passive antibodies may be used to treat people with active COVID-19 in order to help them recover.'"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                }
            ],
            "source": [
                "display(response.strip())"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3 (ipykernel)",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.11.1"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}